package com.niodata.dt.fs.ftp;

import com.google.common.base.Strings;
import com.niodata.dt.fs.DpFtpFileSystem;
import java.io.IOException;
import java.net.Authenticator;
import java.net.ConnectException;
import java.net.InetSocketAddress;
import java.net.PasswordAuthentication;
import java.net.Proxy;
import java.net.Proxy.Type;
import java.net.SocketAddress;
import java.util.Queue;
import java.util.concurrent.ArrayBlockingQueue;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.commons.net.ftp.FTP;
import org.apache.commons.net.ftp.FTPClient;
import org.apache.commons.net.ftp.FTPClientConfig;
import org.apache.commons.net.ftp.FTPReply;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.net.NetUtils;

public class FtpClientPool {

  public static final int DEFAULT_BUFFER_SIZE = 1024 * 1024;
  public static final String FS_FTP_USER_PREFIX = "fs.ftp.user.";
  public static final String FS_FTP_HOST = "fs.ftp.host";
  public static final String FS_FTP_HOST_PORT = "fs.ftp.host.port";
  public static final String FS_FTP_PASSWORD_PREFIX = "fs.ftp.password.";
  private static final Log logger = LogFactory.getLog(FtpClientPool.class);
  private static final int DEFAULT_POOL_SIZE = 10;
  private final Queue<PooledFtpClient> pool;
  private final Configuration config;

  public FtpClientPool(Configuration config) {
    this(DEFAULT_POOL_SIZE, config);
  }

  public FtpClientPool(int maxPoolSize, Configuration config) {
    this.config = config;
    pool = new ArrayBlockingQueue<PooledFtpClient>(maxPoolSize);
  }

  /**
   * Return an active FTPClient.
   */
  public PooledFtpClient getClient() throws IOException {
    PooledFtpClient client = pool.poll();

    if (client == null) {
      client = connection();
    }

    return prepareClient(client);
  }

  private PooledFtpClient connection() throws IOException {
    PooledFtpClient client = new PooledFtpClient(this);
    FTPClientConfig clientConfig = new FTPClientConfig();
    clientConfig.setServerTimeZoneId(config.get(DpFtpFileSystem.FS_FTP_TIMEZONE, "UTC"));
    client.configure(clientConfig);
    String proxyHost = config.get(DpFtpFileSystem.FS_FTP_PROXY_HOST);
    String proxyPort = config.get(DpFtpFileSystem.FS_FTP_PROXY_PORT);
    if (!Strings.isNullOrEmpty(proxyHost)) {
      logger.info("create ftp client with proxy:" + proxyHost);
      SocketAddress addr = new InetSocketAddress(config.get(DpFtpFileSystem.FS_FTP_PROXY_HOST),
            config.getInt(DpFtpFileSystem.FS_FTP_PROXY_PORT, 65535));
      Proxy proxy = new Proxy(Type.SOCKS, addr);
      Authenticator.setDefault(new Authenticator() {
        @Override
        protected PasswordAuthentication getPasswordAuthentication() {
          String proxyUser = config.get(DpFtpFileSystem.FS_FTP_PROXY_USER, null);
          String proxyPassword = config.get(DpFtpFileSystem.FS_FTP_PROXY_PASSWORD, null);
          if ("FTP".equalsIgnoreCase(super.getRequestingProtocol())
                && proxyHost.equals(super.getRequestingHost())
                && proxyPort.equals(super.getRequestingPort())
                && !Strings.isNullOrEmpty(proxyUser)
                && !Strings.isNullOrEmpty(proxyPassword)) {
            return new PasswordAuthentication(proxyUser, proxyPassword.toCharArray());
          }
          return super.getPasswordAuthentication();
        }
      });
      client.setProxy(proxy);
    }
    String host = config.get(FS_FTP_HOST);
    int port = config.getInt(FS_FTP_HOST_PORT, FTP.DEFAULT_PORT);
    String user = config.get(FS_FTP_USER_PREFIX + host);
    String password = config.get(FS_FTP_PASSWORD_PREFIX + host);
    client.connect(host, port);
    logger.info("create ftp client:" + host);
    int reply = client.getReplyCode();
    if (!FTPReply.isPositiveCompletion(reply)) {
      throw NetUtils.wrapException(host, port,
            NetUtils.UNKNOWN_HOST, 0,
            new ConnectException("Server response " + reply));
    } else if (client.login(user, password)) {
      //client.
      client.setFileTransferMode(FTP.BLOCK_TRANSFER_MODE);
      client.setFileType(FTP.BINARY_FILE_TYPE);
      client.setBufferSize(DEFAULT_BUFFER_SIZE);
      client.enterLocalPassiveMode();
      client.setKeepAlive(true);
    } else {
      throw new IOException("Login failed on server - " + host + ", port - "
            + port + " as user '" + user + "'");
    }

    return client;
  }

  /**
   * Prepares ftp client.
   *
   * @param client client
   */
  protected PooledFtpClient prepareClient(PooledFtpClient client) throws IOException {
    return isAlive(client) ? client : getClient();
  }

  protected boolean isAlive(FTPClient client) {
    try {
      if (client.sendNoOp()) {
        return true;
      }
    } catch (IOException e) {
      logger.error("Client [" + client + "] discarded: ", e);
    }

    return false;
  }

  public void releaseClient(PooledFtpClient client) {
    if ((client != null) && isAlive(client) && !pool.offer(client)) {
      try {
        client.trulyLogOut();
        client.trulyDisconnect();
      } catch (Exception e) {
        logger.error("Error disconnecting ftpClient", e);
      }
    }
  }

}
